package org.example.controller;

import jakarta.annotation.Resource;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.stream.Collectors;

@RequestMapping("ollama")
@RestController
public class OllamaController {
    @Resource
    private OllamaChatClient ollamaChatClient;

    @GetMapping("/ai/chat")
    public String getOllamaResponse(@RequestParam String msg) {
        //同步调用 当前页面会卡住 直到所有数据结束才返回
        return ollamaChatClient.call(msg);
    }

    @GetMapping("/ai/stream")
    public Flux<ChatResponse> stream(@RequestParam String msg) {
       //流式返回
        Prompt prompt = new Prompt(new UserMessage(msg));
        Flux<ChatResponse> stream = ollamaChatClient.stream(prompt);
        return stream;
    }
    @GetMapping("/ai/stream1")
    public List<String> stream1(@RequestParam String msg) {
        //后端打印输出内容  只是在后端范围
        Prompt prompt = new Prompt(new UserMessage(msg));
        Flux<ChatResponse> res = ollamaChatClient.stream(prompt);
       List<String> result= res.toStream().map(it->{
            String content = it.getResult().getOutput().getContent();
            System.out.println(content);
            return content;
        }).collect(Collectors.toList());
        return result;
    }

}
